import os
import torch
import argparse
import random
import wave
import python_speech_features as psf
#import torchaudio
import soundfile as sf
import numpy as np
from tqdm import tqdm

parser = argparse.ArgumentParser()
parser.add_argument('-f', '--wavscp', type=str, default=None)
parser.add_argument('-n', '--num_mel_bins', type=int, default=80)
parser.add_argument('-p', '--percent_of_total_sents', type=float, default=1.0)
parser.add_argument('-o', '--outputdir', type=str, default=None)
args = parser.parse_args()


global_mean = 0
global_var = 0
num_frames = 0

random.seed(1234)

def get_var_mean(n, mean1, var1, m, mean2, var2):
    alpha = n / (m+n)
    beta = m / (m+n)
    mean =  alpha * mean1 + beta * mean2
    var = alpha * (var1 + mean1 ** 2) + beta * (var2 + mean2 ** 2) - mean ** 2
    return n+m, mean, var

with open(args.wavscp, 'r', encoding='utf-8') as f:
    lines = f.readlines()
    random.shuffle(lines)

    for i in tqdm(range(int(args.percent_of_total_sents * len(lines)))):
        path = lines[i].strip().split()[-1]
        # wavform, sample_rate = sf.read(path)
        # wavform = torch.FloatTensor(wavform).reshape(1, -1)
        # feature = torchaudio.compliance.kaldi.fbank(wavform, num_mel_bins=args.num_mel_bins, sample_frequency=sample_rate, dither=0.0)
        ob = wave.open(path, 'rb')
        params = ob.getparams()
        nchannels, samplewidth, sample_rate, nframes = params[:4]
        str_data = ob.readframes(nframes)
        wavform = np.fromstring(str_data, dtype=np.short())
        wavform.shape = 1, -1
        wavform = wavform.T
        wavform = torch.tensor(wavform)

        feature_np = psf.base.logfbank(wavform, samplerate=16000, nfilt=self.params['num_mel_bins'])
        feature = torch.tensor(feature_np, dtype=torch.float32)

        var, mean = torch.var_mean(feature, dim=0)
        num_frames, global_mean, global_var = get_var_mean(num_frames, global_mean, global_var, feature.shape[0], mean, var)

global_std = torch.sqrt(global_var)

assert args.outputdir is not None
np.save(os.path.join(args.outputdir, 'global_cmvn.mean'), global_mean.numpy())
np.save(os.path.join(args.outputdir, 'global_cmvn.std'), global_std.numpy())
print('Global Mean and Std are saved as %s' % args.outputdir)

